# lama/community_utils.py

import networkx as nx


def get_neighbors(G, node):
    """获取某个节点在图中的所有邻居"""
    return list(G.neighbors(node))


def is_in_layer(G, node, layer="ppi"):
    """判断节点属于哪一层"""
    return G.nodes[node].get('layer', '') == layer


def expand_community(current_C, current_N, G, z, threshold=0.5):
    """
    根据当前邻居的归属值扩展社区。
    :param current_C: 当前核心社区节点列表
    :param current_N: 当前边界节点列表
    :param G: 融合后的多层网络
    :param z: 节点归属字典 {node_id: membership_value}
    :param threshold: 扩展社区的阈值
    :return: 新加入社区的节点列表
    """
    new_nodes = []
    for node in current_N:
        score = sum(G[u][node]['weight'] * z.get(u, 0.5) for u in G.neighbors(node) if u in current_C)
        if score > threshold:
            new_nodes.append(node)
    return new_nodes


def get_community_subgraph(G, community):
    """返回社区对应的子图"""
    return G.subgraph(community)


def get_local_structure(G, community):
    """提取社区的本地结构特征（如平均度、聚类系数）"""
    subG = get_community_subgraph(G, community)
    avg_degree = sum(dict(subG.degree(weight='weight')).values()) / len(subG)
    density = nx.density(subG)
    clustering = nx.average_clustering(subG)
    return {
        "avg_degree": avg_degree,
        "density": density,
        "clustering": clustering
    }


def compute_node_score(G, node, C, z, delta_w=None):
    """
    计算一个节点的归属得分，用于决定是否将其加入社区。
    :param G: 图对象
    :param node: 待评估节点
    :param C: 当前社区节点集合
    :param z: 毎个节点的归属强度字典 {node: score}
    :param delta_w: 各网络层的权重 {'PPI': 0.6, 'Expression': 0.4}
    :return: 节点的综合得分
    """
    if delta_w is None:
        delta_w = {'PPI': 0.5, 'Expression': 0.5}  # 默认各占一半

    total_score = 0.0
    layer_scores = {}

    # 获取该节点的所有邻居
    neighbors = G[node]

    for neighbor in neighbors:
        if ':' in str(neighbor):
            layer, n_id = neighbor.split(':', 1)
        else:
            layer, n_id = 'unknown', neighbor

        if layer not in layer_scores:
            layer_scores[layer] = 0.0

        if int(n_id) in C:
            layer_scores[layer] += G[node][neighbor]['weight'] * z.get(int(n_id), 0.5)

    # 加权融合不同层的得分
    for layer, score in layer_scores.items():
        weight = delta_w.get(layer, 0.5)
        total_score += weight * score

    return total_score
